import numpy as np
from utils import row_major_to_col_major


if __name__ == "__main__":
    m = 32
    k = 32
    n = 16
    a_filename = "../datas/matrix_a.txt"
    b_filename = "../datas/matrix_b.txt"
    for row_major_a, row_major_b in [
        [False, False],
        [False, True],
        [True, False],
        [True, True],
    ]:
        # A行优先存储、B行优先存储，C都是列优先存储
        if row_major_a and row_major_b:
            A = file_to_np(a_filename, m, k, row_major_a)
            B = file_to_np(b_filename, k, n, row_major_b)
            c_filename = "../datas/matrix_c_a_row_major-b_row_major.txt"
            truth = data_to_mem(c_filename)
            C = np.matmul(A, B)
            # 因为是行优先存储计算出的结果，需要把结果转换为列优先存储
            C = row_major_to_col_major(C.flatten().tolist(), n)
            # 直接比内存中的数据
            np.testing.assert_almost_equal(np.array(truth), np.array(C), decimal=1e-4)
        elif row_major_a and not row_major_b:
            A = file_to_np(a_filename, m, k, row_major_a)
            B = np.array(col_major_to_row_major(data_to_mem(b_filename), n)).reshape(
                k, n
            )
            c_filename = "../datas/matrix_c_a_row_major-b_col_major.txt"
            C = np.matmul(A, B)
            C = row_major_to_col_major(C.flatten().tolist(), n)
            truth = data_to_mem(c_filename)
            np.testing.assert_almost_equal(truth, C, decimal=1e-4)
        elif not row_major_a and row_major_b:
            A = file_to_np(a_filename, m, k, row_major_a)
            # 行优先数据读取
            B = file_to_np(b_filename, k, n, row_major_b)
            # B = row_major_to_col_major(B, n)
            c_filename = "../datas/matrix_c_a_col_major-b_row_major.txt"
            C = np.matmul(A, B)
            C = row_major_to_col_major(C.flatten().tolist(), n)
            truth = data_to_mem(c_filename)
            np.testing.assert_almost_equal(truth, C, decimal=1e-4)
        else:
            A = file_to_np(a_filename, m, k, row_major_a)
            # 行优先数据读取
            # B = file_to_np(b_filename, k, n, not row_major_b)
            B = data_to_mem(b_filename)
            B = col_major_to_row_major(B, n)
            B = np.array(B).reshape(k, n)
            c_filename = "../datas/matrix_c_a_col_major-b_col_major.txt"
            C = np.matmul(A, B)
            C = row_major_to_col_major(C.flatten().tolist(), n)
            truth = data_to_mem(c_filename)
            np.testing.assert_almost_equal(truth, C, decimal=1e-4)
